import random
from wrappers import *
import numpy as np
from matplotlib import pyplot as plt
from nes_py.wrappers import JoypadSpace
import gym_super_mario_bros
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT
import torch
from tqdm import tqdm
from utils import train_off_policy_agent
from DQN import DQN
import utils
from ReplayBuffer import ReplayBuffer

ENV_NAME = 'SuperMarioBros-1-1-v0'
# 对环境进行预处理
env = wrap_environment(ENV_NAME, SIMPLE_MOVEMENT)


# env = gym_super_mario_bros.make('SuperMarioBros-v0')
# env = JoypadSpace(env, SIMPLE_MOVEMENT)

def test_performance(agent, env):
    for i in range(5):
        state = env.reset()
        done = False
        rewards = 0
        while not done:
            env.render()
            action = agent.take_action(state)
            next_state, reward, done, _ = env.step(action)
            rewards += reward
            state = next_state
        print("episode:{} score:{}".format(i, rewards))
    env.close()


if __name__ == '__main__':
    '''
    lr = 2e-3
    num_episodes = 800
    hidden_dim = 128
    gamma = 0.98
    epsilon = 0.01
    # 每10次更新，更新一次目标网络
    target_update = 10
    buffer_size = 10000
    minimal_size = 500
    batch_size = 64
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device(
        "cpu")

    random.seed(0)
    np.random.seed(0)
    env.seed(0)
    torch.manual_seed(0)
    replay_buffer = ReplayBuffer(buffer_size)
    state_dim = env.observation_space.shape[0]

    action_dim = env.action_space.n
    agent = DQN(state_dim, hidden_dim, action_dim, lr, gamma, epsilon,
                target_update, device, isConv=True, input_shape=env.observation_space.shape)

    return_list = []
    for i in range(10):
        with tqdm(total=int(num_episodes / 10), desc='Iteration %d' % i) as pbar:
            for i_episode in range(int(num_episodes / 10)):
                episode_return = 0
                state = env.reset()
                done = False

                while not done:
                    # env.render()
                    action = agent.take_action(state)
                    next_state, reward, done, _ = env.step(action)
                    replay_buffer.add(state, action, reward, next_state, done)
                    state = next_state
                    episode_return += reward
                    # 当buffer数据的数量超过一定值后,才进行Q网络训练
                    if replay_buffer.size() > minimal_size:
                        b_s, b_a, b_r, b_ns, b_d = replay_buffer.sample(batch_size)
                        transition_dict = {
                            'states': b_s,
                            'actions': b_a,
                            'next_states': b_ns,
                            'rewards': b_r,
                            'dones': b_d
                        }
                        agent.update(transition_dict)
                return_list.append(episode_return)
                if (i_episode + 1) % 10 == 0:
                    pbar.set_postfix({
                        'episode':
                            '%d' % (num_episodes / 10 * i + i_episode + 1),
                        'return':
                            '%.3f' % np.mean(return_list[-10:])
                    })
                pbar.update(1)

    # env.close()
    # 保存模型
    print("保存整个模型")
    torch.save(agent, "DQN_agent_mario.pth")

    episodes_list = list(range(len(return_list)))
    plt.plot(episodes_list, return_list)
    plt.xlabel('Episodes')
    plt.ylabel('Returns')
    plt.title('DQN on {}'.format("mario"))
    plt.show()

    mv_return = utils.moving_average(return_list, 9)
    plt.plot(episodes_list, mv_return)
    plt.xlabel('Episodes')
    plt.ylabel('Returns')
    plt.title('DQN on {}'.format("mario"))
    plt.show()

    print("保存模型参数")
    torch.save(agent.q_net.state_dict(), "DQN_agent_parm.pth")
    '''

    agent = torch.load("DQN_agent_mario.pth")
    test_performance(agent, env)
